package partition;

import network.Device;
import network.Network;
import network.Rule;
import verifier.JiffyVerifier;
import verifier.Ports.PersistentPorts;
import verifier.util.Changes;

import java.math.BigInteger;
import java.util.ArrayList;
import java.util.stream.Collectors;

public class SubnetPartition extends Partition {
    private long subnetIp;
    private int prefix;

    public SubnetPartition(String name, Network network, long subnetIp, int prefix) {
        super(name, network);
        this.subnetIp = subnetIp;
        this.prefix = prefix;
    }

    /**
     * Evaluate mini-batch
     * @param M #updates per-mini-batch
     * @return Model update time (us per-update)
     */
    public double test(int M) {
        /**
         Important: each APKeep instance re-labels every device when initialized.
         Given the Device instances are shared, only one APKeep instance can work at the same time.
         */
        this.verifier = new JiffyVerifier(network, new PersistentPorts());
        int subnet = this.verifier.bddEngine.encodeIpv4(BigInteger.valueOf(this.subnetIp), this.prefix);

        int N = 0;
        ArrayList<ArrayList<Rule>> ruless = new ArrayList<>();

        int cnt = 0, r = 0;
        for (Device device : this.getNetwork().getAllDevices()) {
            for (Rule rule : device.getInitialRules().stream().filter(rule -> this.verifier.bddEngine.and(subnet, this.verifier.bddEngine.encodeIpv4(rule.getMatch(), 16)) != 0).collect(Collectors.toSet())) {
                if (cnt % M == 0) {
                    ruless.add(new ArrayList<>());
                    N ++;
                }
                ruless.get(N - 1).add(rule);
                cnt ++;
            }
        }

        System.out.println("# mini-batches " + ruless.size());
        long s = System.nanoTime();
        double tot = 0;
        // ArrayList<Double> times = new ArrayList<>();

        boolean flag1 = false, flag2 = false;
        for (ArrayList<Rule> rules : ruless) {
            if (flag1) System.out.println();
            if (flag1) System.out.println("===========================================");

            if (flag1) System.out.println("#rules in this mini-batch " + rules.size());
            double t0 = System.nanoTime();
            Changes changes = verifier.insertMiniBatch(rules);
            double t1 = System.nanoTime();
            if (flag1) System.out.println("Rule => change " + ((t1 - t0) / rules.size() / 1000) + " us per update");

//            if (flag1) System.out.println("Remaining BDDs: " + (BDDEngine.refCnt - BDDEngine.defCnt));

            if (flag1) System.out.println("# changes (raw) " + changes.aggr0Size());
            if (flag1) System.out.println("# changes (step-1-aggr) " + changes.aggr1Size());
            changes.aggrBDDs();
            if (flag1) System.out.println("# changes (step-2-aggr) " + changes.aggr2Size());

            double t2 = System.nanoTime();
            verifier.update(changes);
            double t3 = System.nanoTime();
            if (flag1) System.out.println("change => ECs " + ((t3 - t2) / rules.size() / 1000) + " us per update");
            // times.add((t3 - t2) / rules.size() / 1000);
            tot += t3 - t2;


            if (flag1) System.out.println("Total #EC: " + verifier.predSize());
        }

        if (flag2) {
            System.out.println("finish " + this.name + ": " + ((System.nanoTime() - s) / cnt / 1000) + " us per-update");
            System.out.println();
//            System.out.println("Total BDD : " + (BDDEngine.tot / cnt / 1000) + " us per-update");
//            System.out.println("BDD time/operation ratio : " + (BDDEngine.tot / BDDEngine.refCnt) + " ns");
//            System.out.println("Remaining BDDs: " + (BDDEngine.refCnt - BDDEngine.defCnt));
            System.out.println("Total #EC: " + verifier.predSize());
            // PrintWriter writer = new PrintWriter("times.log", "UTF-8");
            // for (Double t : times) writer.println(t);
            // writer.close();
            System.out.println("Change => EC: " + (tot / cnt / 1000) + " us per-update");
            System.exit(0);
        }

        return (tot / cnt / 1000);
    }
}
